import psutil
import subprocess
import os
import logging
import threading
import time

logging.basicConfig(filename='./numa.log', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
NUMASTAT_TIME_INTERVAL = 60
TUNING_TIME_INTERVAL = 300


class NUMA_Tuning():
    """ NUMA类,获取NUMA信息,进行核心绑定和内存迁移操作 """
    def __init__(self):
        self.cpu_nums = psutil.cpu_count()
        self.node_count = 0
        self.nodes = {}
        self.processes = []
        self.migration_processes = []
        self.own_pid = os.getpid()
        self.excluded_cmds = ['su', 'bash', 'sleep']
        self.numa_cross = False

        '''启用一个线程，每隔 NUMASTAT_TIME_INTERVAL 秒获取一次numastat信息'''
        thread = threading.Thread(target=self.record_numastat)
        thread.daemon = True  # 将线程设置为守护线程，主程序退出时，线程也会退出
        thread.start()
    
    def tuning(self):
        print("开始NUMA调优, 需要root权限")
        while True:
            self.get_numa_node_info()
            while not self.numa_cross: # 检测是否存在较多跨NUMA内存访问
                self.print_numastat()
                time.sleep(NUMASTAT_TIME_INTERVAL)
            self.print_numastat()
            print("检测到较多跨NUMA内存访问, 开始调优")
            self.get_userspace_processes()
            self.print_userspace_processes()
            self.get_process_numa_stats()
            self.evaluate_migration_candidates()
            self.bind_and_migrate_processes()
            self.numa_cross = False
            time.sleep(TUNING_TIME_INTERVAL) ## 间隔一段时间再次检查

    def parse_numastat(self):
        """ 运行 numastat 命令并记录local_node和other_node的值 """ 
        result = subprocess.run(['numastat'], stdout=subprocess.PIPE)
        stats = result.stdout.decode('utf-8').splitlines()

        # 解析输出以找到local_node和other_node的值
        stats_dict = {}
        for line in stats:
            parts = line.split()
            if parts[0] in ['local_node', 'other_node']:
                # 输出是每个节点一个值，按顺序是 node0, node1, ...
                stats_dict[parts[0]] = list(map(int, parts[1:]))
                
        return stats_dict

    def record_numastat(self):
        previous_stats = self.parse_numastat()
        while True:
            time.sleep(NUMASTAT_TIME_INTERVAL)
            current_stats = self.parse_numastat()
            
            # 计算增量
            deltas = {}
            for key in ['local_node', 'other_node']:
                deltas[key] = [current - prev for current, prev in zip(current_stats[key], previous_stats[key])]

            for i in range(self.node_count): 
                self.nodes[str(i)]['local_node'] = current_stats['local_node'][i]
                self.nodes[str(i)]['other_node'] = current_stats['other_node'][i]
                self.nodes[str(i)]['local_node_deltas'] = deltas['local_node'][i]
                self.nodes[str(i)]['other_node_deltas'] = deltas['other_node'][i]

            # 更新之前的统计数据
            previous_stats = current_stats

    def check_numa_cross(self):
        """ 检查是否存在较多跨NUMA内存访问 """
        for i in range(self.node_count):
            local_node = self.nodes[str(i)]['local_node']
            other_node = self.nodes[str(i)]['other_node']
            local_node_deltas = self.nodes[str(i)]['local_node_deltas']
            other_node_deltas = self.nodes[str(i)]['other_node_deltas']

            if other_node > local_node * 0.1 or other_node_deltas > local_node_deltas * 0.1:
                self.numa_cross = True

    def print_numastat(self):
        for i in range(self.node_count):
            local_node= self.nodes[str(i)]['local_node']
            other_node = self.nodes[str(i)]['other_node']
            local_node_deltas = self.nodes[str(i)]['local_node_deltas']
            other_node_deltas = self.nodes[str(i)]['other_node_deltas']
            print(f"Node {i}: local_node: {local_node}, other_node: {other_node}, local_node_deltas: {local_node_deltas}, other_node_deltas: {other_node_deltas}")
            
    def get_numa_node_info(self):
        """使用 numactl 获取 NUMA 节点信息 """
        result = subprocess.run(['numactl', '--hardware'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        result_str = result.stdout.decode('utf-8')

        # 解析 NUMA 节点数量
        self.node_count = int(result_str.splitlines()[0].split(':')[1].split(' ')[1])

        # 初始化节点信息存储结构
        self.nodes = {str(i): {'cpus': [], 'size': 0, 'free': 0, 'distances': {}, 
                               'local_node': 0, 'other_node': 0, 'local_node_deltas':0, 'other_node_deltas':0} for i in range(self.node_count)}

        # 解析每个节点的 CPU 列表、内存大小和空闲内存
        lines = result_str.splitlines()
        for line in lines:
            if line.startswith('node'):
                parts = line.split()
                node_id = parts[1]

                if 'cpus' in line:
                    self.nodes[node_id]['cpus'] = parts[3:]
                elif 'size' in line:
                    self.nodes[node_id]['size'] = int(parts[3])
                elif 'free' in line:
                    self.nodes[node_id]['free'] = int(parts[3])

            # 解析节点距离
            if line.startswith('node distances:'):
                index = lines.index(line) + 2
                for i in range(self.node_count):
                    distances = lines[index + i].split()[1:]
                    for j in range(self.node_count):
                        self.nodes[str(i)]['distances'][str(j)] = int(distances[j])

    def print_numa_node_info(self):
        print(f"Available NUMA nodes: {self.node_count}")
        for node_id, node_info in self.nodes.items():
            print(f"Node {node_id}:")
            print(f"  CPUs: {', '.join(node_info['cpus'])}")
            print(f"  Size: {node_info['size']} MB")
            print(f"  Free: {node_info['free']} MB")
            print(f"  Distances: {node_info['distances']}")

    def get_userspace_processes(self):
        """ 获取全部用户进程（不对内核进程进行操作）"""
        result = subprocess.run(['ps', '-ef'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        lines = result.stdout.splitlines()

        self.processes = []
        # 跳过标题行
        for line in lines[1:]:
            if line.strip() == "":
                continue

            parts = line.split()
            uid, pid, ppid, c, stime, tty, time = parts[:7]
            cmd = " ".join(parts[7:])
            
            # 忽略内核线程
            if (cmd.startswith('[') and cmd.endswith(']')) or int(pid) == self.own_pid or int(ppid) == self.own_pid:
                continue
            
            if any(cmd.startswith(excluded_cmd) for excluded_cmd in self.excluded_cmds):
                continue
            
            self.processes.append({
                'PID': pid,
                'CMD': cmd,
            })
    
    def print_userspace_processes(self):
        for process in self.processes:
            print(f"PID: {process['PID']}, CMD: {process['CMD']}")


    def get_process_numa_stats(self):
        """ 获取每个进程的 NUMA 统计信息 """
        for process in self.processes:
            pid = process['PID']
            result = subprocess.run(['numastat', '-p', pid], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            lines = result.stdout.splitlines()

            total_line = lines[-1] # 最后一行是总额
            node_usage = total_line.split()[1:]
            node_stats = {}

            for i in range(self.node_count):
                node_stats[f"Node_{i}"] = float(node_usage[i])
            node_stats['Total'] = float(node_usage[-1])
            process['numa_stats'] = node_stats


    def print_process_numa_stats(self):
        for process in self.processes:
            print(f"PID {process['PID']}: \n\tCMD: {process['CMD']} \n\tNUMA Stats: {process['numa_stats']}")

    def evaluate_migration_candidates(self, threshold_MB=20):
        """ 筛选出需要进行内存迁移的进程，然后进行排序 """
        self.migration_processes = []
        for process in self.processes:
            pid = process['PID']
            numa_stats = process['numa_stats']
            node_memories = [float(numa_stats[f'Node_{i}']) for i in range(self.node_count)]
            max_memory = max(node_memories)
            main_node = str(node_memories.index(max_memory))
            total_memory = float(numa_stats['Total'])
            other_memory = total_memory - max_memory

            # 跳过内存小的进程
            if max_memory < threshold_MB:
                continue  # 跳过这个进程，不考虑迁移

            self.migration_processes.append({
                'PID': pid,
                'numastat': numa_stats,
                'Total Memory': total_memory,
                'Other Memory': other_memory,
                'Main Node': main_node
            })
            
        # 筛选过后的进程，按总内存大小进行排序
        self.migration_processes.sort(key=lambda x: x['Total Memory'], reverse=True)

        for process in self.migration_processes:
            print(f"PID: {process['PID']}, Total Memory: {process['Total Memory']} MB, Other Memory: {process['Other Memory']} MB, Main Node: {process['Main Node']}")

    def print_migration_candidates(self):
        for process in self.migration_processes:
            print(f"PID: {process['PID']}, Total Memory: {process['Total Memory']} MB, Other Memory: {process['Other Memory']} MB, Main Node: {process['Main Node']}")

    def bind_and_migrate_processes(self):
        """ 绑定核心，迁移进程到目标节点 """
        num_processes_to_migrate = 1 # 对前n个进程进行迁移
        migrated_processes = 0
        for process in self.migration_processes:
            if migrated_processes >= num_processes_to_migrate:
                break
            
            # 尝试迁移内存
            self.update_node_info()
            if process['Other Memory'] < self.nodes[process['Main Node']]['free']: # 检查目标结点是否有内存可迁移
                self.bind_and_migrate_one_process(pid=process['PID'], target_node=process['Main Node'])
                migrated_processes += 1

    def update_node_info(self):
        """ 更新self.nodes的可用内存信息 """
        result = subprocess.run(['numactl', '--hardware'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        result_str = result.stdout.decode('utf-8')

        # 解析每个节点的 CPU 列表、内存大小和空闲内存
        lines = result_str.splitlines()
        for line in lines:
            if line.startswith('node'):
                parts = line.split()
                node_id = parts[1]
                if 'free' in line:
                    self.nodes[node_id]['free'] = int(parts[3])


    def bind_and_migrate_one_process(self, pid, target_node):
        """ 绑定核心，迁移进程到目标节点 """
        logging.info(f"尝试绑定进程 {pid} 到结点 {target_node}")
        # 绑定核心
        cpu_list = ','.join(self.nodes[str(target_node)]['cpus'])
        subprocess.run(['taskset', '-pc', f'{cpu_list}', f'{pid}'])
        logging.info(f"taskset -pc {cpu_list} {pid}")

        # 迁移内存
        for node in range(self.node_count):
            node = str(node)
            if node == target_node:
                continue
            else:
                res = subprocess.run(['migratepages', pid, node, target_node])
                logging.info(f"migratepages {pid} {node} {target_node}")
    

if __name__ == '__main__':
    numa = NUMA_Tuning()
    numa.tuning()
